
import pyvista as pv
import numpy as np

mesh = pv.read("u2.obj")
mesh = mesh.triangulate()
#mesh.plot()

from sketchc import *
window = Window((1000, 800))

if "ball":
    x, = mesh. faces.shape
    l = mesh.faces.reshape((x // 4, 4))
    l = l[:, 1:4]
    l = l.flatten()
    vertices = mesh.points[l]
    normals = mesh.point_normals[l]
    texture_coords = mesh.point_data["材质1"][l]


    ball = Objt("tex1.png")
    ball.color = V3(0.5, 0.5, 0.5)
    ball.color = V3(0, 0, 0)
    ball.color = V3(1, 1, 1)

    (vertices, normals, texture_coords) >> ball
    ball.u(window.ctx)
    ball >> window
    

    
    

if "shader":
    喷漆shader = {
        "vertex_shader": """
#version 330 core
in vec3 aPos;
in vec3 aNormal;
in vec2 aTex;

uniform mat4 modelMat;
uniform mat4 viewMat;
uniform mat4 projMat;
uniform mat4 template_trans;

out vec4 wPos;

void main()
{
    wPos = modelMat * vec4(aPos + aNormal * 0, 1.0);
    wPos = inverse(template_trans) * wPos;
    vec4 wwPos = modelMat * vec4(aTex.x, 1- aTex.y, 0, 1.0);
    gl_Position = projMat * viewMat * wwPos;
}
        """,
        "fragment_shader": """
#version 330 core
out vec4 FragColor;

uniform vec2 template_size;
uniform float d1;
uniform float d2;
uniform sampler2D texture;

uniform vec3 lightPos;
uniform vec3 viewPos;
uniform vec3 objectColor;
uniform int get_normal;

in vec4 wPos;

float sign2(float x)
{
    return (sign(x) + 1) / 2;
}

void main()
{
    float useless = lightPos.x * viewPos.x * objectColor.x * float(get_normal) * 0;
    vec3 tag = vec3(useless, 0, -d1 - d2);
    vec3 de = wPos.xyz - tag;
    vec2 r = de.xy / de.z * (d1 + d2);
    r = r / template_size * 2;
    float x = r.x - r.y;
    float y = r.x + r.y;
    float k = sign2(pow(2, 0.5) - abs(x) - abs(y)) * sign2(de.z - d2);
    
    vec4 result = texture2D(texture, (r + vec2(1, 1)) / 2);
    
    FragColor = mix(vec4(0, 0, 0, 0), result, k);
}
        """
    }

if "喷漆":
    board = Objt("tex3.png")
    board.prog = 喷漆shader
    board.color = V3(1, 1, 1)
    
    (vertices[::-1], normals[::-1], texture_coords[::-1]) >> board
    
    board.u(window.ctx)
    board.prog["d1"].value = 1.5
    board.prog["d2"].value = 0
    board.prog["template_size"].write(np.array([2, 2], dtype="f4"))    
    
    # ~ if "frame" and False:
    if "frame":
        
        boardframe = Objt("tex3.png")
        boardframe.mode = moderngl.LINES
        window.ctx.line_width = 2
        boardframe.prog = 喷漆shader
        boardframe.color = V3(1, 1, 1)
        
        a, b, c = vertices[:,0], vertices[:,1], vertices[:,2]
        d, e, f = normals[:,0], normals[:,1], normals[:,2]
        g, h = texture_coords[:,0], texture_coords[:,1]
        x = np.stack([a, b, c, d, e, f, g, h], axis=1)
        x0 = x[::3, :]
        x1 = x[1::3, :]
        x2 = x[2::3, :]
        xx = np.stack([x0, x1, x1, x2, x2, x0], axis=1)
        m, n, o = xx.shape
        xx = xx.reshape((m * n, o))
        
        
        
        
        (xx[::-1, 0:3], xx[::-1, 3:6], xx[::-1, 6:8]) >> boardframe
        
        boardframe.u(window.ctx)
        boardframe.prog["d1"].value = 1.5
        boardframe.prog["d2"].value = 0
        boardframe.prog["template_size"].write(np.array([2, 2], dtype="f4"))    
    
    
    
    base = Obj()
    temp = Obj()
    temp.t.iwalk(rot=vrot(U3.up, -0.35))
    temp.t.iwalk(U3(0, -1.7, 0))
    temp.t.connect(base.t)
    temp2 = Obj()
    temp2.t.iwalk(rot=vrot(U3.up, 0.35))
    temp2.t.iwalk(U3(0, -1.7, 0))
    temp2.t.connect(base.t)
    

    
    
    # ~ board.prog["template_trans"].write(temp.m())
    # ~ boardframe.prog["template_trans"].write(temp.m())


if "white":
    white = Objt("white.png")
    white.prog = 喷漆shader
    white.color = V3(1, 1, 1)
    
    (vertices[::-1], normals[::-1], texture_coords[::-1]) >> white
    
    white.u(window.ctx)
    white.prog["d1"].value = 1.5
    white.prog["d2"].value = 0
    white.prog["template_size"].write(np.array([0.8, 1.6], dtype="f4"))    
    
    
    t3 = Turtle(rot=vrot(U3.up, -0.35))
    t3.connect(base.t)
    t4 = Turtle(rot=vrot(U3.up, 0.35))
    t4.connect(base.t)
    t5 = Turtle(rot=vrot(U3.lt, 1.2))       # 1.2睁开，0.6闭上
    t5.parent = t3
    t6 = Turtle(rot=vrot(U3.lt, 1.2))   
    t6.parent = t4
    
    wobj = Obj()

camera = Camera(U3(0, -5, 0))
camera >> window

walker = Walker(camera)
walker >> window

fbo2 = EasyFbo((1024, 1024), window.ctx, quality=True)
board_camera = Camera(U3(0, -10, 0))
board_camera.proj = Matrix44.orthogonal_projection(0, 1, 0, 1, 0, 1024)

window.ctx.polygon_offset = -1.0, -1.0

@window.update_callback.append
def f(window, dt):
    # ~ base.t.walk(rot=vrot(U3.up, dt * 2))
    
    fbo2.use()
    window.ctx.enable(moderngl.BLEND)
    window.ctx.disable(moderngl.DEPTH_TEST)
    fbo2.ctx.clear(1, 1, 1, 1)
    
    # 开始喷漆
    board.prog["template_trans"].write(temp.m())
    board.draw(board_camera)
    boardframe.prog["template_trans"].write(temp.m())
    boardframe.draw(board_camera)
    board.prog["template_trans"].write(temp2.m())
    board.draw(board_camera)
    boardframe.prog["template_trans"].write(temp2.m())
    boardframe.draw(board_camera)
    
    wobj.t = t5
    white.prog["template_trans"].write(wobj.m())
    white.draw(board_camera)
    wobj.t = t6
    white.prog["template_trans"].write(wobj.m())
    white.draw(board_camera)
    
    #结束喷漆
    
    window.ctx.disable(moderngl.BLEND)
    window.ctx.enable(moderngl.DEPTH_TEST)
    data = fbo2.texture.read()
    mat = np.frombuffer(data, dtype="f4").reshape((1024, 1024, 4))
    mat = (mat * 255).astype("uint8")
    ball.texture.write(mat)
    
    window.ctx.screen.use()
    
from PIL import Image

@window.key_down_callback.append
def f(window, e):
    if e.key == pygame.K_RETURN:
        data = fbo2.texture.read()
        mat = np.frombuffer(data, dtype="f4").reshape((1024, 1024, 4))
        mat = (mat * 255).astype("uint8")
        img = Image.fromarray(mat)
        img.save("texout.png")
        print("saved")
        
@window.key_down_callback.append
def f(window, e):
    if e.key == pygame.K_f:
        def update1(window, dt):
            # ~ print("update1")
            x = update1.acc
            if x < 0.2:
                update1.acc += dt
            else:
                window.update_callback.remove(update1)
            
            def shape(x):
                if 0 <= x <= 1:
                    return 1 - abs(1 - 2 * x)
                else:
                    return 0
                
            th = 1.2  - shape(x / 0.2) * 0.6 
            # ~ print(th)
            rot = vrot(U3.lt, th)
            t5.rot = rot
            t6.rot = rot
        
        update1.acc = 0
        window.update_callback.append(update1)
        


clock = pygame.time.Clock()
while True:
    dt = clock.tick(30)
    window.update(dt)
    for e in pygame.event.get():
        if e.type == pygame.QUIT:
            exit()
        window.dispatch(e)
    window.render()
    pygame.display.flip()
    
    






